import warnings

import pandas as pd
import torch

from .distance_map import *
from .evaluate import *
from model import LayerNormNet, CNN, RNN, LSTM
from .utils import *


def warn(*args, **kwargs):
    pass
warnings.warn = warn

def infer_pvalue(p_value = 1e-5, nk_random = 20,
                 train_data='/state/partition/wzzheng/clean/data/train_valid_split/split100/split100_train_split_2',
                 test_data='/state/partition/wzzheng/clean/data/datasets/new',
                 esm_dir='/state/partition/wzzheng/clean/data/train_valid_split/split100/esm_data',
                 model_name='cnn',
                 model_path='/state/partition/wzzheng/clean/pretrained/split100_triplet.pth',
                 results_dir='/state/partition/wzzheng/clean/results'):
    seed_everything()
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:0" if use_cuda else "cpu")
    dtype = torch.float32

    # load ec id dictionary
    id_ec_train, ec_id_dict_train = get_ec_id_dict(train_data + '.csv')
    id_ec_test, _ = get_ec_id_dict(test_data + '.csv')

    # load checkpoints
    # NOTE: change this to LayerNormNet(512, 256, device, dtype) 
    # and rebuild with [python build.py install]
    # if inferencing on model trained with supconH loss
    if model_name == 'cnn':
        model = CNN(1280, 512, 128)
    elif model_name == 'rnn':
        model = RNN(1280, 512, 128)
    elif model_name == 'lstm':
        model = LSTM(1280, 512, 128)
    elif model_name == 'layernorm':
        model = LayerNormNet(512, 128, device, dtype)
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint)
    model.to(device, dtype)
    model.eval()

    # # load training EC cluster center embeddings 
    emb_train = model(esm_embedding(ec_id_dict_train, device, dtype, esm_dir=esm_dir))

    # # calculate distance map and p-value for test data
    emb_test = model_embedding_test(id_ec_test, model, device, dtype, esm_dir=esm_dir)
    eval_dist = get_dist_map_test(emb_train, emb_test, ec_id_dict_train, id_ec_test, device, dtype)
    eval_df = pd.DataFrame.from_dict(eval_dist)
    rand_nk_ids, rand_nk_emb_train = random_nk_model(
        id_ec_train, ec_id_dict_train, emb_train, n=nk_random, weighted=True)
    random_nk_dist_map = get_random_nk_dist_map(
        emb_train, rand_nk_emb_train, ec_id_dict_train, rand_nk_ids, device, dtype)
    
    ensure_dirs(results_dir)
    test_name = os.path.basename(test_data)
    out_filename = os.path.join(results_dir, model_name + test_name + '_pvalue.csv')
    print(f'>>> Writing p-value choices to {out_filename}')
    write_pvalue_choices(eval_df, out_filename, random_nk_dist_map, p_value=p_value)

    pred_label = get_pred_labels(out_filename)
    pred_probs = get_pred_probs(out_filename)
    true_label, all_label = get_true_labels(test_data + '.csv')
    pre, rec, f1, roc, acc = get_eval_metrics(
        pred_label, pred_probs, true_label, all_label)
    print(f'############ EC calling results using random '
    f'chosen {nk_random}k samples ############')
    print('-' * 75)
    print(f'>>> total samples: {len(true_label)} | total ec: {len(all_label)} \n'
        f'>>> precision: {pre:.3} | recall: {rec:.3}'
        f'| F1: {f1:.3} | AUC: {roc:.3} | Accuracy: {acc:.3} \n')
    print('-' * 75)  